#define USE_SYS_LIBC
#define USE_SYS_MSVCRT
#include "ntlib.hpp"
#include <Windns.h>
#include <stdlib.h>

NTSYSAPI
    LONG
    NTAPI
    RtlIpv4StringToAddressW (
    _In_ PCWSTR S,
    _In_ BOOLEAN Strict,
    _Out_ LPCWSTR *Terminator,
    _Out_ struct in_addr *Addr
    );
#pragma comment(linker,"/entry:DllEntry")
#pragma comment(lib,"Dnsapi.lib")
#pragma comment(lib,"ws2_32.lib")

//template struct __declspec(dllexport) ntlib<0>;
WCHAR module[MAX_PATH];
HMODULE curr_module;
WCHAR proc[MAX_PATH];

WCHAR** filters = 0;
WCHAR adns[] = L"8.8.8.8\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0";
struct tuple{
	WCHAR* q;
	WCHAR* r;
} *gsub = 0;
extern "C" BOOL WINAPI DnsFlushResolverCache();

static HKEY key;
static HANDLE hEvent;

struct shared_struct 
{
    bool init;
    bool is_debug;
};
shared_struct* _shared = 0;
shared_struct _default = {false,false};
static HANDLE shared_mapping;
int process();
void setup();
void cleanup();

struct DNS_SERVER
{
    WCHAR* domain,*host,*dot1;
    DWORD ds1;
    DWORD ds2;
    QWORD us1;
    PVOID nls;
    DWORD ds3;
    DWORD flag1;
    QWORD flag2;
    WCHAR *guid ,*domain2 ;
    QWORD us3;
    WCHAR *dot2;
    PDNS_ADDR_ARRAY z_array;
};
struct DNS_Q {
    wchar_t* name;
    wchar_t* name2;
    wchar_t* name3;
    DWORD query_type;
    DWORD flags;
    QWORD us[14];
    PDNS_RECORDW records;
    DWORD us2[2];
    //PIP4_ARRAY arry;  
    DNS_SERVER* arry;
};

void* init_shared(wchar_t* dll, size_t size)
{
    void md5s(void *msg, size_t len, wchar_t out[]);
    wchar_t md5[64] ={0};
    md5s(dll,wcslen(dll)*sizeof(dll[0]),md5);
    if (md5[0])
    {
        return ntlib::shared_section(size,md5,shared_mapping);
    }
    return 0;
}

#define dprint(...) if (_shared->is_debug) ntlib::debug_print(__VA_ARGS__);
#define dprintw(...)  if(_shared->is_debug) ntlib::debug_printw(__VA_ARGS__);

using namespace ntlib;
typedef PVOID (WINAPI *Query_Main_t)(void*);
typedef PDNS_RECORD (WINAPI *DnsAllocateRecord_t)(DWORD);
static DnsAllocateRecord_t DnsAllocateRecord;
hook_info hook;
Query_Main_t real_query;

ULONG tcp_only(PWCHAR name)
{
    if (filters) {
        WCHAR** p = filters;
        while(*p != 0) {
            if (wcsstr(name,*p))
                return Inet_Addr(adns);
            p++;
        }
    }
    return INADDR_NONE;
}
PWCHAR subsitute(PWCHAR name)
{
    if (gsub) {
        WCHAR* match = name;
        bool matched = false;
again:
        tuple* a = &gsub[0];
        while (a->q && a->r) {
            if (wcsstr(match,a->q)) {
                matched = true;
                match = a->r;
                goto again;
            }
            a++;
        }

        if (matched) {
            return match;
        }
    }
    return name;
}
PVOID WINAPI Query_Main(DNS_Q* query)
{
    if(_shared->is_debug) ntlib::debug_printw(L"Query_Main: %s(type:%X flag:%X arry:%I64X)",query->name,query->query_type,query->flags,query->arry);

    if (query->query_type != DNS_TYPE_A && query->query_type != DNS_TYPE_AAAA)
        return real_query(query);

    auto fdns = tcp_only(query->name);
    if (fdns != INADDR_NONE)           // Use TCP dns
    {
        query->flags |= DNS_QUERY_USE_TCP_ONLY;
        sockaddr_in * addr = (sockaddr_in *)&(query->arry->z_array[0].AddrArray->MaxSa);
        addr->sin_family = AF_INET;
        addr->sin_port = 0;
        addr->sin_addr.S_un.S_addr = fdns;

    }
    else
    {

        PWCHAR fname = subsitute(query->name);
        PWCHAR oname = query->name;

        auto address = Inet_Addr(fname);
        if (address != INADDR_NONE)  // A raw address
        {
            PDNS_RECORDW record = (PDNS_RECORDW)DnsAllocateRecord(sizeof(DNS_A_DATA));
            record->pName = oname;
            record->Data.A.IpAddress = address;
            record->wType = DNS_TYPE_A;
            record->dwTtl = 86400;
            record->Flags.S.Section = DNSREC_ANSWER;
            record->Flags.S.CharSet = DnsCharSetUnicode;
            //record->Flags.S.Reserved = 0X20; //iF there is more records afterwards
            query->records = record;
            return 0;
        }

        query->name = fname;                  // A name

        auto freedns = tcp_only(fname);
        if (freedns != INADDR_NONE)           // Use TCP dns
        {
            query->flags |= DNS_QUERY_USE_TCP_ONLY;
            sockaddr_in * addr = (sockaddr_in *)&(query->arry->z_array[0].AddrArray->MaxSa);
            addr->sin_family = AF_INET;
            addr->sin_port = 0;
            addr->sin_addr.S_un.S_addr = freedns;
        }
    }

    PVOID res =  real_query(query);
    return res;
}
int hook_query()
{
    HMODULE dns = GetModuleHandleW(L"dnsapi.dll");
    DnsAllocateRecord = (DnsAllocateRecord_t)GetProcAddress(dns,"DnsAllocateRecord");
    if (hook_proc(GetProcAddress(dns,"Query_Main"), Query_Main,&real_query,15,&hook))
    {
        dprint("hook success: %d",real_query);
        return 1;
    }
    return 0;
}
BOOL WINAPI DllEntry(_In_  HINSTANCE hinstDLL, _In_  DWORD reason, _In_  LPVOID lpvReserved)
{
    if (reason == DLL_PROCESS_ATTACH) {
        DisableThreadLibraryCalls(hinstDLL);
        curr_module = hinstDLL;
        GetModuleFileNameW(hinstDLL,module,sizeof(module)/sizeof(module[0]));
        GetModuleFileNameW(0,proc,sizeof(proc)/sizeof(proc[0]));
        (void*&)_shared = init_shared(module,sizeof(*_shared));
        if (!_shared->init)
        {
            memcpy(_shared,&_default,sizeof(_default));
            _shared->init = true;
        }
        dprintw(L"Load %s in %s",module,proc);
        if (GetModuleHandleW(L"dnsrslvr.dll")) 
        {
            GetModuleHandleExW(0,module,&hinstDLL);   // hold on to this dll
            if (hook_query())
            {
                setup();
            }
        }
    }
    else if (reason == DLL_PROCESS_DETACH) {
        if (unhook_proc(&hook)) 
        {
            cleanup();
        }
        dprintw(L"Free %s in %s",module,proc);
        ntlib::release_section(_shared,shared_mapping);
    }
    return TRUE;
}
#ifndef _WIN64
#pragma comment(linker,"/export:queryW=_Query@16")
#else
#pragma comment(linker,"/export:queryW=Query")
#endif
extern "C" void WINAPI Query(HWND hwndStub, HINSTANCE hAppInstance,LPTSTR lpszCmdLine,int nCmdShow)
{
    PDNS_RECORD set = 0;
    PVOID re = _alloca(1000);
    IP4_ARRAY ss;
    PIP4_ARRAY pss= 0;
    int argc = 0;
    wchar_t** argv = argvw(lpszCmdLine,&argc);
    if (argc >=4){
        ss.AddrCount = 1;
        ss.AddrArray[0] = inet_addr("8.8.8.8");
        pss = &ss;
    }
    DNS_STATUS  stat = DnsQuery_W(argv[0],wcstoul(argv[1],0,0),wcstoul(argv[2],0,0),pss,&set,0);
    while(set){
        debug_print("%S %s ttl:%d\n",set->pName, inet_ntoa( *((in_addr*)&(set->Data.A.IpAddress))),set->dwTtl );
        set = set->pNext;
    }
}
#ifndef _WIN64
#pragma comment(linker,"/export:configW=_config1@16")
#else
#pragma comment(linker,"/export:configW=config1")
#endif
extern "C" void WINAPI config1(HWND hwndStub, HINSTANCE hAppInstance,LPTSTR lpszCmdLine,int nCmdShow)
{
    int argc = 0;
    wchar_t** argv= argvw(lpszCmdLine,&argc);
    if (lstrcmpiW(argv[0],L"-debug") == 0)
    {
        ntlib::debug_print("dns debug: %d", _shared->is_debug);
        _shared->is_debug = !_shared->is_debug;
    }
    else 
    {
        setup();
    }
}

DWORD __stdcall freeme(void*) {
    FreeLibraryAndExitThread(curr_module,0);
    return 1;
}

DWORD getdnscache() {
    SC_HANDLE s = OpenSCManagerW(0,0,GENERIC_READ);
    SC_HANDLE sv = OpenServiceW(s,L"dnscache",SERVICE_QUERY_STATUS);
    SERVICE_STATUS_PROCESS status = {0};
    DWORD r;
    QueryServiceStatusEx(sv,SC_STATUS_PROCESS_INFO ,(LPBYTE )&status,sizeof(status),&r);
    return status.dwProcessId;
}
#ifndef _WIN64
#pragma comment(linker,"/export:injW=_inject@16")
#else
#pragma comment(linker,"/export:injW=inject")
#endif
extern "C" void WINAPI inject(HWND,HINSTANCE,LPWSTR, int )
{
    DWORD pid = getdnscache();
    if (pid) {
        PVOID p = GetProcAddress(GetModuleHandleW(L"kernel32.dll"),"GetModuleHandleW");
        DWORD code = 1;
        if (call_remote_par(pid,(tproc_t)p,module,wcslen(module)*sizeof(WCHAR),&code) && code == 0) {
            if (inject_me(pid,module) == 0) {
                dprint("Failed injection %d",GetLastError());
            }
        }
        else {
            dprint("Cannot call GetModuleHandleW");
        }
    }
    else {
        dprint("Cannot find dnscache");
    }
}
#ifndef _WIN64
#pragma comment(linker,"/export:freeW=_free_@16")
#else
#pragma comment(linker,"/export:freeW=free_")
#endif
extern "C" void WINAPI free_(HWND,HINSTANCE,LPWSTR, int )
{
    DWORD pid = getdnscache();
    if (pid) {
        call_remote(pid,freeme,0);
    }
}


DWORD WINAPI wfunc(PVOID) {
	while(1){
		WaitForSingleObject(hEvent,INFINITE);
		process();
		RegNotifyChangeKeyValue(key,0,REG_NOTIFY_CHANGE_LAST_SET,hEvent,1);
	}
}
static HANDLE bthread;
void setup() {
	RegOpenKeyExW(HKEY_LOCAL_MACHINE,L"SYSTEM\\CurrentControlSet\\Control\\ip",0,KEY_QUERY_VALUE|KEY_NOTIFY,&key);
	hEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
	process();
	RegNotifyChangeKeyValue(key,0,REG_NOTIFY_CHANGE_LAST_SET,hEvent,1);
	bthread = CreateThread(0,0,wfunc,0,0,0);
}
void free_gsub()
{
    auto *it = gsub;
    while (it && it->q)
    {
        free(it->q);
        free(it->r);
        it++;
    }
    free(gsub);
}
void cleanup() {
	if (key == 0) return;
	TerminateThread(bthread,0);
	CloseHandle(bthread);
	CloseHandle(hEvent);
	RegCloseKey(key);
	free_gsub();
	free(filters);
}

void sort_subs()
{
    auto * pivot = gsub, *it = gsub;
    while (it && it->q)
    {
        if (it->q[0] == L'!')
        {
            auto temp = *pivot;
            *pivot = *it;
            *it = temp;
            memmove(pivot->q,pivot->q+1,wcslen(pivot->q)*2);
            pivot++;
        }
        it++;
    }
}

int process() {

	DWORD ret = 0,type =0;
	DWORD size = 0;
	DWORD count = 0, namelen = 0, valuelen = 0;
	RegQueryInfoKeyW(key,0,0,0,0,0,0,&count,&namelen,&valuelen,0,0);
	namelen += 1;
	WCHAR* name = (WCHAR*)_alloca(namelen*2);
	LPBYTE  value = (LPBYTE)_alloca(valuelen);
	tuple* subs = (tuple*)malloc(count*sizeof(tuple));
	int ii = 0;
	for (unsigned  q = 0; q < count; q++) {
		DWORD nl = namelen, vl = valuelen;
		ret = RegEnumValueW(key,q,name,&nl,0,0,value,&vl);
		if (ret == ERROR_SUCCESS && wcscmp(name,L"afilter") != 0) {
			subs[ii].q = (WCHAR*)malloc(nl*2+2);
			wcscpy(subs[ii].q,name);
			subs[ii].r =  (WCHAR*)malloc(vl+2);
			wcsncpy(subs[ii++].r,(WCHAR*)value,vl/2);
		}
	}
	subs[ii].q=0;
    free_gsub();
	gsub = subs;
    sort_subs();
	WCHAR* buf = (WCHAR*)malloc(valuelen);
	size = valuelen;
	ret = RegQueryValueExW(key,L"afilter",0,&type,(PBYTE)buf,&size);
	if (ret != ERROR_SUCCESS)
		return -1;
	WCHAR** f = (WCHAR**)malloc(size/2*sizeof(char*));
	size_t i = 0;
	int j = 0;
	while (i < size) {
		size_t len = wcslen(&buf[i]);
		if (len > 0) {
			f[j++] = &buf[i];
			i += len + 1;
		}
		else {
			break;
		}
	}        
	f[j] = 0;
	if (filters && filters[0]) free(filters[0]);
	free(filters);
	filters = f;

    size = sizeof(adns);
    ret = RegQueryValueExW(key,L"_adns",0,&type,(PBYTE)adns,&size);

	DnsFlushResolverCache();
	return 0;
}
